# 将可视化pkl结果文件与原json文件建立映射关系，产生json文件
import json
import pickle
import os
import sys
from vis_class import asm_node

# 部分硬编码路径
vis_pkl_path = r'preprocessed_data/addon_test/addon_vis/vis_output/'
json_output_path = r'preprocessed_data/addon_test/addon_vis/vis_json/' 

def auto_vis():
    for fname in os.listdir(vis_pkl_path):
        pkl_file = open(os.path.join(vis_pkl_path, fname), 'rb')
        vis_output = pickle.load(pkl_file)
        # 每个vis文件分析结果都单独保存为两个部分，0和1
        for filename in vis_output.keys():
            output_file_name = filename.replace('/', '@')
            print(output_file_name)
            visinfo = vis_output[filename]
            json_file = open(filename, 'r')
            json_load = json.load(json_file)
            json_nodes = json_load["nodes"]

            # 构建每个汇编节点的信息，建立汇编节点地址到汇编节点信息的映射字典
            dict_asmaddr_2_asminfo = {}
            for func in json_nodes:
                idx = 0
                for asm in func["disasm"]:
                    node = asm_node(func["id"], func["addr"], func["function_name"], idx, asm["address"], asm["mnemonic"], asm  ["op_str"], 0, 0)
                    idx += 1
                    dict_asmaddr_2_asminfo[asm["address"]] = node

            # 基于pkl决策影响和json汇编信息建立映射关系，并存到文件里，这里分0/1
            output_ls_0 = []
            for vis0 in visinfo["0"]:
                target_addr = vis0[1]
                if target_addr == -1:
                    continue
                dict_asmaddr_2_asminfo[target_addr].ctb_0 = float(vis0[0][0])
                dict_asmaddr_2_asminfo[target_addr].ctb_1 = float(vis0[0][1])
                output_ls_0.append(dict_asmaddr_2_asminfo[target_addr].to_json())

            output_ls_1 = []
            for vis1 in visinfo["1"]:
                target_addr = vis1[1]
                if target_addr == -1:
                    continue
                dict_asmaddr_2_asminfo[target_addr].ctb_0 = float(vis1[0][0])
                dict_asmaddr_2_asminfo[target_addr].ctb_1 = float(vis1[0][1])
                output_ls_1.append(dict_asmaddr_2_asminfo[target_addr].to_json())

            output_file_path = os.path.join(json_output_path, output_file_name)
            output_f = open(output_file_path, 'w')
            json.dump({'0':output_ls_0,'1':output_ls_1}, output_f)
            output_f.close()




if __name__ == '__main__':    
    auto_vis()